const vscode = require('vscode');

/**
 * MCP工具调用结果
 */
class McpToolCallResult {
    constructor() {
        this.call = null;
        this.result = null;
        this.error = null;
    }
}

/**
 * MCP助手类，用于在AI代理中集成MCP功能
 */
class McpAssistant {
    constructor(mcpService) {
        this.mcpService = mcpService;
        this.pendingApprovals = new Map();
        this.approvalCount = 0;
    }

    /**
     * 获取可用的MCP工具列表
     */
    getAvailableTools() {
        try {
            return this.mcpService.getAvailableTools();
        } catch (error) {
            console.error('获取MCP工具列表失败:', error);
            return [];
        }
    }

    /**
     * 为AI助手构建MCP工具的功能描述
     */
    buildMcpToolsDescription() {
        const tools = this.getAvailableTools();
        
        if (!tools || tools.length === 0) {
            return '';
        }
        
        // 按服务器分组工具
        const toolsByServer = {};
        
        for (const tool of tools) {
            if (!toolsByServer[tool.serverName]) {
                toolsByServer[tool.serverName] = [];
            }
            toolsByServer[tool.serverName].push(tool);
        }
        
        // 构建工具描述
        let description = `你有以下MCP工具可以使用:\n\n`;
        
        for (const serverName in toolsByServer) {
            description += `## ${serverName} 服务器\n\n`;
            
            for (const tool of toolsByServer[serverName]) {
                description += `### ${tool.toolName}\n`;
                if (tool.description) {
                    description += `${tool.description}\n`;
                }
                
                // 添加参数信息展示
                if (tool.properties) {
                    description += `**参数:**\n`;
                    try {
                        // 直接处理properties对象，它包含参数定义
                        const params = tool.properties;
                        
                        // 遍历所有参数
                        for (const [paramName, paramDetails] of Object.entries(params)) {
                            // 检查参数是否必填
                            const required = tool.required && tool.required.includes(paramName);
                            
                            // 添加参数描述
                            description += `- \`${paramName}\`${required ? ' (必填)' : ''}: ${paramDetails.type || '未知类型'} - ${paramDetails.description || '无描述'}\n`;
                        }
                        
                        description += '\n';
                    } catch (e) {
                        console.error('处理工具参数信息失败:', e);
                        description += `参数信息解析失败\n\n`;
                    }
                }
                
                description += `使用方式: <mcp-tool server="${serverName}" tool="${tool.toolName}" params={...参数...}></mcp-tool>\n\n`;
            }
        }
        
        // 添加使用说明
        description += `
## MCP工具使用说明

你可以在回复中使用MCP工具完成各种任务，语法如下:

\`\`\`
<mcp-tool server="服务器名称" tool="工具名称" params={...参数...}></mcp-tool>
\`\`\`

其中:
- server: MCP服务器名称
- tool: 要使用的工具名称
- params: 工具所需的参数，格式为JSON对象，请参考各工具的参数说明

MCP工具使用规则（必须严格遵守）:

1. 标签必须正确闭合：每个<mcp-tool>必须有对应的</mcp-tool>闭合标签
2. 一次回复中只调用一个工具：每次回复最多只能包含一个工具调用
3. 等待工具响应：调用工具后必须立即结束当前回复，等待工具执行结果后再在下一次调用中继续
4. 不要提前完成：在完全完成用户初始请求前，不要使用attemptCompletion标签
5. 逐步执行：一次只能执行一个步骤，等待结果后再决定下一步操作

例如：

\`\`\`
<mcp-tool server="openai" tool="retrieveText" params={"query": "查询内容"}></mcp-tool>
\`\`\`

错误示例（不要这样做）:
1. 不要在同一回复中调用多个工具
2. 不要在未收到工具响应前继续回复
3. 不要在未完成用户请求前使用attemptCompletion
4. 不要使用未闭合的标签

工具调用可能需要用户确认，或者在配置了自动批准的情况下直接执行。
`;
        
        return description;
    }

    /**
     * 处理AI助手的响应，提取并执行MCP工具调用
     * @param {string} response AI助手的原始响应
     * @param {function} requestApproval 请求用户批准的回调函数
     * @returns {Promise<{processedResponse: string, toolResults: Array}>} 处理后的响应和工具调用结果
     */
    async processMcpResponse(
        response,
        requestApproval
    ) {
        const toolResults = [];
        let processedResponse = response;
        
        try {
            console.log(`[MCP处理] 开始处理MCP响应`);
            // 提取MCP工具调用
            const regex = /<mcp-tool\s+server="([^"]+)"\s+tool="([^"]+)"\s+params=(\{[^}]+\})><\/mcp-tool>/g;
            const matches = [...response.matchAll(regex)];
            
            console.log(`[MCP处理] 提取到 ${matches.length} 个MCP工具调用`);
            
            if (matches.length === 0) {
                return { processedResponse, toolResults };
            }
            
            // 处理每个工具调用
            for (const match of matches) {
                const [fullMatch, serverName, toolName, paramsStr] = match;
                let params = {};
                
                try {
                    // 解析参数
                    params = JSON.parse(paramsStr);
                    
                    // 创建工具调用参数
                    const toolCall = {
                        serverName,
                        toolName,
                        params
                    };
                    
                    console.log(`[MCP处理] 处理MCP工具调用: ${serverName}.${toolName}`, params);
                    
                    // 检查工具是否需要用户批准
                    const needsApproval = !this.mcpService.isToolAutoApproved(serverName, toolName);
                    console.log(`[MCP处理] 工具是否需要批准: ${needsApproval}`);
                    
                    if (needsApproval && typeof requestApproval === 'function') {
                        // 请求用户批准
                        console.log(`[MCP处理] 请求用户批准: ${serverName}.${toolName}`);
                        const result = await new Promise((resolve, reject) => {
                            const approvalId = `approval_${this.approvalCount++}`;
                            
                            // 设置超时
                            const timeoutId = setTimeout(() => {
                                reject(new Error('用户批准超时'));
                            }, 60000); // 60秒超时
                            
                            this.pendingApprovals.set(approvalId, { 
                                params: toolCall, 
                                resolve: (result) => {
                                    console.log(`[MCP处理] 用户批准回调 - 批准结果:`, result);
                                    clearTimeout(timeoutId);
                                    resolve(result);
                                }, 
                                reject: (err) => {
                                    console.log(`[MCP处理] 用户批准回调 - 拒绝结果:`, err);
                                    clearTimeout(timeoutId);
                                    reject(err);
                                }
                            });
                            
                            // 调用回调函数请求用户批准，传递服务器名称
                            requestApproval(toolName, params, (approved) => {
                                console.log(`[MCP处理] 用户${approved ? '批准' : '拒绝'}了工具调用: ${serverName}.${toolName}`);
                                if (approved) {
                                    this.executeTool(toolCall)
                                        .then(result => {
                                            console.log(`[MCP处理] 工具执行成功:`, result);
                                            resolve(result);
                                        })
                                        .catch(error => {
                                            console.log(`[MCP处理] 工具执行失败:`, error);
                                            reject(error);
                                        })
                                        .finally(() => {
                                            // 从挂起的批准中删除
                                            this.pendingApprovals.delete(approvalId);
                                            clearTimeout(timeoutId);
                                        });
                                } else {
                                    console.log(`[MCP处理] 用户拒绝，准备拒绝回调`);
                                    reject(new Error('用户拒绝了工具调用'));
                                    
                                    // 从挂起的批准中删除
                                    this.pendingApprovals.delete(approvalId);
                                    clearTimeout(timeoutId);
                                }
                            }, serverName);
                        }).catch(error => {
                            console.error(`[MCP处理] 工具批准/执行失败: ${error.message}`);
                            return { error };
                        });
                        
                        // 添加结果
                        toolResults.push({
                            call: toolCall,
                            result: result.error ? undefined : result,
                            error: result.error
                        });
                    } else {
                        // 自动批准，直接执行
                        console.log(`自动批准执行: ${serverName}.${toolName}`);
                        try {
                            const result = await this.executeTool(toolCall);
                            console.log(`工具执行结果:`, result);
                            toolResults.push({
                                call: toolCall,
                                result
                            });
                        } catch (error) {
                            console.error(`工具执行失败: ${error.message}`);
                            toolResults.push({
                                call: toolCall,
                                error
                            });
                        }
                    }
                } catch (error) {
                    console.error(`处理MCP工具调用失败: ${error.message}`);
                    toolResults.push({
                        call: {
                            serverName,
                            toolName,
                            params
                        },
                        error
                    });
                }
            }
            
            // 替换响应中的工具调用为结果
            for (let i = 0; i < matches.length; i++) {
                const [fullMatch] = matches[i];
                const result = toolResults[i];
                
                let replacement = '';
                if (!result) {
                    replacement = `<mcp-error>工具调用处理失败</mcp-error>`;
                } else if (result.error) {
                    replacement = `<mcp-error server="${result.call.serverName}" tool="${result.call.toolName}">错误: ${result.error instanceof Error ? result.error.message : String(result.error)}</mcp-error>`;
                } else {
                    let resultStr = '';
                    try {
                        resultStr = typeof result.result === 'object' 
                            ? JSON.stringify(result.result, null, 2) 
                            : String(result.result);
                    } catch (error) {
                        resultStr = '[无法序列化结果]';
                    }
                    
                    replacement = `<mcp-result server="${result.call.serverName}" tool="${result.call.toolName}">${resultStr}</mcp-result>`;
                }
                
                // 使用替换
                processedResponse = processedResponse.replace(fullMatch, replacement);
            }
            
            return { processedResponse, toolResults };
        } catch (error) {
            console.error(`处理MCP响应出错:`, error);
            return { 
                processedResponse: response + `\n<mcp-error>处理MCP工具调用出错: ${error.message}</mcp-error>`, 
                toolResults 
            };
        }
    }

    /**
     * 执行MCP工具调用
     * @param {object} toolCall 工具调用参数
     * @returns {Promise<any>} 工具执行结果
     */
    async executeTool(toolCall) {
        try {
            return await this.mcpService.callTool(toolCall);
        } catch (error) {
            console.error(`执行MCP工具调用失败: ${error instanceof Error ? error.message : String(error)}`);
            throw error;
        }
    }
}

module.exports = { McpAssistant }; 